# -*- coding: utf-8 -*-
# =============================================================================
#         Desc: 用于删除指定目录或指定文件的Page Cache
#       Author: GGA
#        Email:
#     HomePage:
#      Version: 1.0.1
#   LastChange: 2020-12-20
#      History:
# =============================================================================

import os
import sys
import math
import datetime
import argparse
import time

from ctypes.util import find_library
from ctypes import (
    c_int,
    c_size_t,
    c_ubyte,
    c_void_p,
    get_errno,
    cdll,
    POINTER,
    cast
)
from mmap import (
    MAP_SHARED,
    PROT_READ,
    PAGESIZE
)

# 默认时间格式
default_datetime_format = "%Y-%m-%d %H:%M:%S"


class FilePageCacheInfo(object):
    def __init__(self, file_path, cache_page_size, cache_page_count, file_page_count):
        self.file_path = file_path
        self.cache_page_size = cache_page_size
        self.cache_page_count = cache_page_count
        self.file_page_count = file_page_count

    @property
    def cache_page_mb(self):
        return round(self.cache_page_size * self.cache_page_count / 1024 / 1024, 2)

    @property
    def file_page_mb(self):
        return round(self.cache_page_size * self.file_page_count / 1024 / 1024, 2)


class PrintHelper(object):
    is_debug = False

    @classmethod
    def get_current_time_string(cls):
        return datetime.datetime.now().strftime(default_datetime_format)

    @classmethod
    def print_info(cls, message):
        print("[{}] {}".format(
            cls.get_current_time_string(),
            message
        ))

    @classmethod
    def print_data(cls, message):
        print(message)

    @classmethod
    def print_debug(cls, message):
        if cls.is_debug:
            print("[{}] {}".format(
                cls.get_current_time_string(),
                message
            ))


class PageCacheMapper(object):
    """
    抄自pgpagecache包
    """

    def __init__(self):
        self.MAP_FAILED = c_void_p(-1)
        self.c_off_t = self.get_c_s_size()
        libc = self.get_lib_c()
        self.mmap = libc.mmap
        self.mmap.restype = c_void_p
        self.mmap.argtypes = [c_void_p, c_size_t, c_int, c_int, self.c_off_t]
        self.munmap = libc.munmap
        self.munmap.restype = c_void_p
        self.munmap.argtypes = [c_void_p, c_size_t]
        self.mincore = libc.mincore
        self.mincore.argtypes = [c_void_p, c_size_t, POINTER(c_ubyte)]
        self.mincore.restype = c_int

    @classmethod
    def get_lib_c(cls):
        libclib = find_library("c")
        if libclib is None:
            raise OSError("Could not load libc dynamic library")
        libc = cdll.LoadLibrary(libclib)
        return libc

    @classmethod
    def get_c_s_size(cls):
        try:
            from ctypes import c_ssize_t
        except ImportError:
            from ctypes import c_longlong
            is_64bits = sys.maxsize > 2 ** 32
            c_ssize_t = c_longlong if is_64bits else c_int
        return c_ssize_t

    def get_page_cache(self, filename):
        with open(filename, "r") as fd:
            stat = os.fstat(fd.fileno())
            filesize = stat.st_size
            if filesize == 0:
                return PAGESIZE, 0, 0
            faddr = self.mmap(0, filesize, PROT_READ, MAP_SHARED, fd.fileno(),
                              0)
            if faddr == self.MAP_FAILED:
                fd.close()
                PrintHelper.print_info(
                    "Failed to mmap %s (errno: %d)" % (filename, get_errno())
                )
                return None, None, None

            vec_size = int((filesize + PAGESIZE - 1) / PAGESIZE)
            vec = (c_ubyte * vec_size)()
            rv = self.mincore(faddr, filesize, cast(vec, POINTER(c_ubyte)))
            if rv == -1:
                fd.close()
                PrintHelper.print_info(
                    "mincore failed: 0x%x, 0x%x: %d" % (
                        faddr, filesize,
                        get_errno())
                )
                return None, None, None
            cached_pages = 0
            for pageidx in range(len(vec)):
                if vec[pageidx] & 1:
                    cached_pages += 1
            # Clean up
            fd.close()
            if faddr:
                self.munmap(faddr, filesize)
            del stat
            del vec
            # Return cached pages, total pages, ratio
            total_pages = math.ceil(float(filesize) / PAGESIZE)
            return PAGESIZE, cached_pages, int(total_pages)


class PageCacheHelper(object):

    @classmethod
    def __get_dir_files(cls, dir_path):
        """
        获取特定目录下所有文件
        :param dir_path:
        :return:
        """
        file_list = []
        for root, dirs, files in os.walk(dir_path):
            for file in files:
                file_list.append(os.path.join(root, file))
        return file_list

    @classmethod
    def __drop_file_page_cache(cls, file_path, sleep_seconds):
        """
        删除指定文件的Page Cache
        :param file_path:
        :param sleep_seconds:
        :return:
        """
        try:
            PrintHelper.print_info("开始清理文件{}的Page Cache".format(file_path))
            pc_info = cls.__get_file_page_cache(file_path=file_path)
            if pc_info is None:
                PrintHelper.print_info("未能获取到文件{}的Page Cache信息".format(file_path))
                return
            with open(file_path, "rb") as fh:
                file_length = os.path.getsize(filename=file_path)
                os.posix_fadvise(fh.fileno(), 0, file_length, os.POSIX_FADV_DONTNEED)
                PrintHelper.print_info("完成清理文件{}的Page Cache".format(file_path))
            sleep_time = round(sleep_seconds * pc_info.cache_page_mb / 1024, 3)
            sleep_info = "文件{}的Page Cache为{}mb, 休眠{}ms。".format(
                file_path,
                pc_info.cache_page_mb,
                sleep_time * 1000
            )
            PrintHelper.print_info(sleep_info)
            time.sleep(sleep_time)
        except Exception as ex:
            PrintHelper.print_info("清理文件{}的Page Cache失败，异常：{}".format(file_path, str(ex)))

    @classmethod
    def __drop_dir_page_cache(cls, dir_path, sleep_seconds, file_extend_name):
        """
        删除特定目录下的Page Cache
        :param dir_path:
        :return:
        """
        file_list = cls.__get_dir_files(dir_path=dir_path)
        for file_path in file_list:
            if cls.check_file_extend_name(
                    file_path=file_path,
                    file_extend_name=file_extend_name
            ):
                cls.__drop_file_page_cache(file_path=file_path, sleep_seconds=sleep_seconds)

    @classmethod
    def __get_file_page_cache(cls, file_path):
        """
        获取指定文件的Page Cache信息
        :param
            file_path:
            page_cache_mapper:
        :return:
        返回数据：
        返回FilePageCacheInfo对象
        当返回None时，标识获取PageCache失败
        """
        try:
            PrintHelper.print_info("开始收集{}文件的PageCache信息".format(file_path))
            page_cache_mapper = PageCacheMapper()
            cache_page_size, cache_page_count, file_page_count = page_cache_mapper.get_page_cache(filename=file_path)
            return FilePageCacheInfo(
                file_path=file_path, cache_page_size=cache_page_size,
                cache_page_count=cache_page_count, file_page_count=file_page_count)
        except Exception as ex:
            PrintHelper.print_info("获取文件{}的Page Cache出错，异常：{}".format(file_path, str(ex)))
            return None

    @classmethod
    def __get_dir_page_cache(cls, dir_path, file_extend_name):
        """
        获取特定目录的Page Cache
        :param dir_path:
        :return:
        返回该该目录下所有文件的Page Cache例表
        """
        cache_info_list = []
        file_list = cls.__get_dir_files(dir_path=dir_path)
        for file_path in file_list:
            if cls.check_file_extend_name(
                    file_path=file_path,
                    file_extend_name=file_extend_name
            ):
                cache_info = cls.__get_file_page_cache(
                    file_path=file_path)
                if cache_info is not None:
                    cache_info_list.append(cache_info)
        return cache_info_list

    @classmethod
    def drop_page_cache(cls, path, sleep_seconds, file_extend_name):
        """
        删除特定目录或特定文件的Page Cache
        :param
            path:
            sleep_seconds:
        :return: void
        """
        if os.path.isfile(path):
            if cls.check_file_extend_name(
                    file_path=path,
                    file_extend_name=file_extend_name):
                cls.__drop_file_page_cache(file_path=path, sleep_seconds=sleep_seconds)
        else:
            cls.__drop_dir_page_cache(
                dir_path=path,
                sleep_seconds=sleep_seconds,
                file_extend_name=file_extend_name
            )

    @classmethod
    def get_page_cache(cls, path, file_extend_name):
        """
        获取特定目录下或特定文件的Page Cache
        :param
            path:
            file_extend_name:
        :return:
        返回特定目录或特定文件的Page Cache例表
        """
        cache_info_list = []
        if os.path.isfile(path):
            is_match_file = cls.check_file_extend_name(
                file_path=path,
                file_extend_name=file_extend_name
            )
            if is_match_file:
                cache_info = cls.__get_file_page_cache(
                    file_path=path
                )
                cache_info_list.append(cache_info)
        else:
            cache_info_list = cls.__get_dir_page_cache(
                dir_path=path,
                file_extend_name=file_extend_name
            )
        return cache_info_list

    @classmethod
    def check_file_extend_name(cls, file_path, file_extend_name):
        if file_extend_name == "*":
            return True
        file_name = os.path.basename(file_path)
        if file_name.endswith(file_extend_name):
            return True
        return False

    @classmethod
    def show_cache_info(cls, path, file_extend_name):
        cache_info_list = cls.get_page_cache(path=path, file_extend_name=file_extend_name)
        total_file_page_count = 0
        total_cache_page_count = 0
        total_file_page_mb = 0
        total_cache_page_mb = 0
        new_cache_info_list = []
        for cache_info in cache_info_list:
            total_file_page_count += cache_info.file_page_count
            total_cache_page_count += cache_info.cache_page_count
            total_file_page_mb += cache_info.file_page_mb
            total_cache_page_mb += cache_info.cache_page_mb
            new_cache_info_list.append(cache_info)
        total_file_page_mb = round(total_file_page_mb, 2)
        total_cache_page_mb = round(total_cache_page_mb, 2)
        new_cache_info_list.sort(key=lambda x: x.cache_page_mb, reverse=False)
        if total_cache_page_mb == 0:
            total_cache_page_mb = 1
        for cache_info in new_cache_info_list:
            cache_rate = round(100 * cache_info.cache_page_mb / total_cache_page_mb, 2)
            info_mess = """ 缓存比率: {}%, 缓存页数: {}/{}, 缓存大小: {}MB/{} MB, 文件: {}""".format(
                cache_rate,
                cache_info.cache_page_count,
                cache_info.file_page_count,
                cache_info.cache_page_mb,
                cache_info.file_page_mb,
                cache_info.file_path
            )
            PrintHelper.print_data(info_mess)
        info_mess = """ 总缓存信息: 缓存页数: {}/{}, 缓存大小: {}MB/{}MB """.format(
            total_cache_page_count, total_file_page_count,
            total_cache_page_mb, total_file_page_mb
        )
        PrintHelper.print_data(info_mess)


class PageCacheChecker(object):
    @classmethod
    def get_parser(cls):
        parser = argparse.ArgumentParser(
            description='查看或清理指定文件(目录)使用的Page Cache',
            add_help=False)
        connect_setting = parser.add_argument_group('参数信息')
        parser.add_argument(
            '--help',
            dest='help',
            action='store_true',
            help='获取帮助信息',
            default=False
        )
        connect_setting.add_argument(
            '--work_path',
            dest='work_path',
            type=str,
            default=None,
            help='要操作的文件路径或目录路径'
        )
        connect_setting.add_argument(
            '--show_detail',
            dest='show_detail',
            action='store_true',
            default=True,
            help='显示指定目录或指定文件的Page Cache详细信息。'
        )
        connect_setting.add_argument(
            '--drop_cache',
            dest='drop_cache',
            action='store_true',
            default=False,
            help='删除指定目录或指定文件的Page Cache。'
        )
        connect_setting.add_argument(
            '--no_confirm',
            dest='no_confirm',
            action='store_true',
            default=False,
            help='设置清理缓存时无需确认,默认值为False,需要确认'
        )
        connect_setting.add_argument(
            '--sleep_seconds',
            dest='sleep_seconds',
            type=float,
            default=1,
            help='清理1 GB Page Cache的休眠时间(秒),默认休眠1秒'
        )
        connect_setting.add_argument(
            '--file_extend_name',
            dest='file_extend_name',
            type=str,
            default="*",
            help='文件后缀名，默认为所有文件'
        )
        return parser

    @classmethod
    def get_user_choose_option(cls, input_options, input_message):
        while_flag = True
        choose_option = None
        while while_flag:
            PrintHelper.print_data(input_message)
            str_input = input("")
            for input_option in input_options:
                if str(str_input).strip() == input_option:
                    choose_option = input_option
                    while_flag = False
        return choose_option

    @classmethod
    def parse_args(cls, command_args):
        need_print_help = False if command_args else True
        parser = cls.get_parser()
        args = parser.parse_args(command_args)
        if args.help or need_print_help:
            parser.print_help()
            sys.exit(1)
        if args.work_path is None:
            PrintHelper.print_info("请输入要查看|清理的文件路径或目录路径")
            parser.print_help()
            sys.exit(1)
        args.work_path = str(args.work_path).strip()
        if (args.work_path == "") or (not os.path.exists(args.work_path)):
            PrintHelper.print_info("请输入要查看|清理的文件路径或目录路径")
            parser.print_help()
            sys.exit(1)
        return args

    @classmethod
    def drop_cache_confirm(cls, work_path):
        input_options = ['yes', 'no']
        input_message = """
    确认要清理{}下的PageCache吗？
    请输入yes确认或输入no退出，yes/no?
    """.format(work_path)
        user_option = cls.get_user_choose_option(
            input_options=input_options,
            input_message=input_message)
        if user_option == "no":
            return False
        else:
            return True

    @classmethod
    def check_cache(cls, command_args):
        args = cls.parse_args(command_args)
        PrintHelper.print_info("开始执行...")
        work_path = args.work_path
        file_extend_name = args.file_extend_name
        sleep_seconds = args.sleep_seconds
        PageCacheHelper.show_cache_info(
            path=work_path,
            file_extend_name=file_extend_name
        )
        if (args.drop_cache == 1) and (args.no_confirm or cls.drop_cache_confirm(work_path)):
            PageCacheHelper.drop_page_cache(
                path=work_path,
                file_extend_name=file_extend_name,
                sleep_seconds=sleep_seconds)
            PageCacheHelper.show_cache_info(
                path=work_path,
                file_extend_name=file_extend_name
            )
        PrintHelper.print_info("完成执行...")


def main():
    PrintHelper.is_debug = True
    PageCacheChecker.check_cache(sys.argv[1:])


if __name__ == '__main__':
    main()
